#ifndef MYNET_H_
#include <torch/torch.h>
#include <vector>
#include <iostream>

// ==================== STN 模块 ======================
struct STNModuleImpl : torch::nn::Module {
    // 定位网络，用于提取全局信息后回归仿射变换参数
    torch::nn::Sequential localization{ nullptr };
    // 全连接层回归仿射矩阵参数，输出 6 个数（对应 2×3 仿射矩阵）
    torch::nn::Sequential fc_loc{ nullptr };

    STNModuleImpl() {
        localization = torch::nn::Sequential(
            // 第1层卷积：输入 1 通道，输出 8 个特征图，核大小 7×7
            torch::nn::Conv2d(torch::nn::Conv2dOptions(1, 8, 7)),
            torch::nn::MaxPool2d(torch::nn::MaxPool2dOptions(2).stride(2)),
            torch::nn::ReLU(true),
            // 第2层卷积：8→10，核大小 5×5
            torch::nn::Conv2d(torch::nn::Conv2dOptions(8, 10, 5)),
            torch::nn::MaxPool2d(torch::nn::MaxPool2dOptions(2).stride(2)),
            torch::nn::ReLU(true)
        );
        register_module("localization", localization);

        // 定位网络输出尺寸为 [10, 3, 3] 展平后为 90
        fc_loc = torch::nn::Sequential(
            torch::nn::Linear(10 * 3 * 3, 32),
            torch::nn::ReLU(true),
            torch::nn::Linear(32, 6)
        );
        register_module("fc_loc", fc_loc);

        // 将 fc_loc 的最后一层初始化为恒等变换 [[1, 0, 0],[0, 1, 0]]
        torch::Tensor identity = torch::tensor({ 1.0, 0.0, 0.0, 0.0, 1.0, 0.0 });
        for (auto& module : fc_loc->modules(/*include_self=*/false)) {
            if (auto* linear = dynamic_cast<torch::nn::LinearImpl*>(module.get())) {
                torch::nn::init::constant_(linear->weight, 0);
                torch::nn::init::constant_(linear->bias, 0);
            }
        }
        if (auto* linear = dynamic_cast<torch::nn::LinearImpl*>((*(fc_loc))[2].get())) {
            torch::nn::init::constant_(linear->weight, 0);
            linear->bias.data().copy_(identity);
        }
    }

    torch::Tensor forward(torch::Tensor x) {
        // x: [B, 1, 28, 28]
        auto xs = localization->forward(x);
        xs = xs.view({ xs.size(0), -1 });          // [B, 90]
        auto theta = fc_loc->forward(xs);          // [B, 6]
        theta = theta.view({ -1, 2, 3 });             // [B, 2, 3]
        // 生成采样网格
        auto grid = torch::nn::functional::affine_grid(theta, x.sizes(), /*align_corners=*/false);
        // 利用采样网格对输入图像进行仿射变换
        auto x_transformed = torch::nn::functional::grid_sample(x, grid,
            torch::nn::functional::GridSampleFuncOptions().align_corners(false));
        return x_transformed;
    }
};
TORCH_MODULE(STNModule);

// ==================== NonLocal Attention 模块 ======================
struct NonLocalBlockImpl : torch::nn::Module {
    torch::nn::Conv2d theta{ nullptr }, phi{ nullptr }, g{ nullptr }, out_conv{ nullptr };

    NonLocalBlockImpl(int channels) {
        theta = register_module("theta", torch::nn::Conv2d(torch::nn::Conv2dOptions(channels, channels / 2, 1)));
        phi = register_module("phi", torch::nn::Conv2d(torch::nn::Conv2dOptions(channels, channels / 2, 1)));
        g = register_module("g", torch::nn::Conv2d(torch::nn::Conv2dOptions(channels, channels / 2, 1)));
        out_conv = register_module("out_conv", torch::nn::Conv2d(torch::nn::Conv2dOptions(channels / 2, channels, 1)));
    }

    torch::Tensor forward(torch::Tensor x) {
        auto batch = x.size(0);
        auto channels = x.size(1);
        auto height = x.size(2);
        auto width = x.size(3);

        auto theta_x = theta->forward(x).view({ batch, -1, height * width });
        auto phi_x = phi->forward(x).view({ batch, -1, height * width });
        auto f = torch::bmm(theta_x.transpose(1, 2), phi_x);  // [B, H×W, H×W]
        f = torch::softmax(f, -1);
        auto g_x = g->forward(x).view({ batch, -1, height * width });
        auto y = torch::bmm(g_x, f.transpose(1, 2));
        y = y.view({ batch, -1, height, width });
        y = out_conv->forward(y);
        return x + y;
    }
};
TORCH_MODULE(NonLocalBlock);

// ==================== 残差块（Residual Block） ======================
struct ResidualBlockImpl : torch::nn::Module {
    torch::nn::Conv2d conv1{ nullptr }, conv2{ nullptr };
    torch::nn::BatchNorm2d bn1{ nullptr }, bn2{ nullptr };

    // 这里 channels 表示输入与输出的通道数
    ResidualBlockImpl(int channels) {
        conv1 = register_module("conv1", torch::nn::Conv2d(
            torch::nn::Conv2dOptions(channels, channels, 3).padding(1)));
        bn1 = register_module("bn1", torch::nn::BatchNorm2d(channels));
        conv2 = register_module("conv2", torch::nn::Conv2d(
            torch::nn::Conv2dOptions(channels, channels, 3).padding(1)));
        bn2 = register_module("bn2", torch::nn::BatchNorm2d(channels));
    }

    torch::Tensor forward(torch::Tensor x) {
        auto residual = x;  // 保存原始输入作为残差
        auto out = conv1->forward(x);
        out = bn1->forward(out);
        out = torch::relu(out);
        out = conv2->forward(out);
        out = bn2->forward(out);
        out += residual;  // 与输入进行逐元素相加
        return torch::relu(out);
    }
};
TORCH_MODULE(ResidualBlock);

// ==================== 主网络 ======================
struct Net : torch::nn::Module {
    STNModule stn{ nullptr };     // STN 模块用于自适应放缩小目标
    torch::nn::Sequential base_conv1{ nullptr }, base_conv2{ nullptr }, base_conv3{ nullptr };
    torch::nn::Conv2d merge_conv{ nullptr };
    ResidualBlock residual_block{ nullptr };  // 新增残差块
    NonLocalBlock non_local{ nullptr };
    torch::nn::Sequential fc_common{ nullptr };
    torch::nn::Linear cls_head{ nullptr };

    Net() {
        // STN 模块：对输入图像做仿射变换，提高小目标的识别能力
        stn = STNModule();
        register_module("stn", stn);

        // base_conv1：1→32，3×3 卷积（padding=1），BN + ReLU，2×2 最大池化 → 输出 [B, 32, 14, 14]
        base_conv1 = torch::nn::Sequential(
            torch::nn::Conv2d(torch::nn::Conv2dOptions(1, 32, 3).padding(1)),
            torch::nn::BatchNorm2d(32),
            torch::nn::ReLU(),
            torch::nn::MaxPool2d(2)
        );
        register_module("base_conv1", base_conv1);

        // base_conv2：32→64，3×3 卷积（padding=1），BN + ReLU，2×2 最大池化 → 输出 [B, 64, 7, 7]
        base_conv2 = torch::nn::Sequential(
            torch::nn::Conv2d(torch::nn::Conv2dOptions(32, 64, 3).padding(1)),
            torch::nn::BatchNorm2d(64),
            torch::nn::ReLU(),
            torch::nn::MaxPool2d(2)
        );
        register_module("base_conv2", base_conv2);

        // base_conv3：64→128，3×3 卷积（padding=1），BN + ReLU，（尺寸保持 7×7）→ 输出 [B, 128, 7, 7]
        base_conv3 = torch::nn::Sequential(
            torch::nn::Conv2d(torch::nn::Conv2dOptions(64, 128, 3).padding(1)),
            torch::nn::BatchNorm2d(128),
            torch::nn::ReLU()
        );
        register_module("base_conv3", base_conv3);

        // 多尺度特征融合：将 base_conv1 的输出 ([B, 32, 14, 14]) 与上采样后的 base_conv3 输出 ([B, 128, 14, 14]) 拼接
        // 拼接后通道数为 32+128=160，再通过 1×1 卷积将特征融合降维为 128 通道
        merge_conv = torch::nn::Conv2d(torch::nn::Conv2dOptions(32 + 128, 128, 1));
        register_module("merge_conv", merge_conv);

        // 残差块：对融合后的特征（128 通道）进行深度特征提取，增加模型的鲁棒性
        residual_block = ResidualBlock(128);
        register_module("residual_block", residual_block);

        // 非局部注意力模块：增强全局特征表达，输入输出均为 [B, 128, 14, 14]
        non_local = NonLocalBlock(128);
        register_module("non_local", non_local);

        // 全连接层：将 128×14×14 的特征展平后映射到 256 维
        fc_common = torch::nn::Sequential(
            torch::nn::Linear(128 * 14 * 14, 256),
            torch::nn::ReLU(),
            torch::nn::Dropout(0.5)
        );
        register_module("fc_common", fc_common);

        // 分类头：将 256 维特征映射到 10 类（数字 0～9）的得分输出
        cls_head = torch::nn::Linear(256, 10);
        register_module("cls_head", cls_head);
    }

    torch::Tensor forward(torch::Tensor x) {
        // 1. STN 模块：对输入图像进行自适应放缩，使较小目标放大后便于识别（输出尺寸保持 [B, 1, 28, 28]）
        x = stn->forward(x);
        // 2. 基础特征提取：base_conv1 输出 [B, 32, 14, 14]
        auto f1 = base_conv1->forward(x);
        // 3. 通过 base_conv2 提取特征，得到 [B, 64, 7, 7]
        auto f2 = base_conv2->forward(f1);
        // 4. 通过 base_conv3 得到深层特征 [B, 128, 7, 7]
        auto f3 = base_conv3->forward(f2);
        // 5. 将 f3 上采样至 14×14
        auto f3_upsampled = torch::nn::functional::interpolate(
            f3,
            torch::nn::functional::InterpolateFuncOptions().scale_factor(std::vector<double>{2, 2})
            .mode(torch::kNearest)
        );
        // 6. 多尺度融合：拼接 f1 ([B, 32, 14, 14]) 与上采样 f3 ([B, 128, 14, 14]) 得到 [B, 160, 14, 14]
        auto fused = torch::cat({ f1, f3_upsampled }, 1);
        // 7. 融合特征经 1×1 卷积降维至 128 通道
        fused = merge_conv->forward(fused);
        // 8. 残差块进一步提取特征，输出仍为 [B, 128, 14, 14]
        fused = residual_block->forward(fused);
        // 9. 通过非局部注意力模块增强全局信息，输出 [B, 128, 14, 14]
        fused = non_local->forward(fused);
        // 10. 将特征展平为向量 [B, 128*14*14]
        auto flattened = fused.view({ fused.size(0), -1 });
        // 11. 全连接层将向量映射到 256 维, 输出 [B, 256]
        auto features = fc_common->forward(flattened);
        // 12. 分类头输出 10 个类别得分 [B, 10]
        auto cls = cls_head->forward(features);
        return cls;
    }
};
/*
下面给出一个修改后的网络代码示例，重点在于增加一个空间变换网络（Spatial Transformer Network, STN）模块。这个模块能够自适应地调整输入图像的尺度和位置，特别是对于较小的目标，可以自动对输入区域进行放大，从而使后续的特征提取模块能获得足够分辨率的信息做分类。整个网络的主要流程如下：

STN 模块   - 使用一个轻量级的定位网络（localization network）提取输入图像中的局部特征，利用全连接层回归出一个仿射变换参数（包含缩放、平移等）。   - 利用仿射变换生成采样网格，对输入图像执行放缩操作。   - 这样，小目标在被放大后进入后续网络，可以更好地保留细节信息。

后续分类网络   - 保留原有的多尺度特征提取与融合：    • base_conv1 提取较高分辨率信息（14×14，32通道）    • base_conv2 与 base_conv3 提取更深层语义（7×7，最终得到 128 通道）    • 将 base_conv3 上采样到 14×14后与 base_conv1 拼接，经过 1×1 卷积融合为 128 通道特征    • 加入非局部注意力模块增强全局特征   - 展平后通过全连接层提取紧凑的全局特征，最后分类头映射到 10 个类别（数字 0~9）。

下面就是完整的修改版 C++ 代码，注意代码中 STN 的实现参考了 PyTorch 官方教程，同时对后续网络保持了多尺度融合和非局部注意力模块的设计：
*/
/*struct Net : torch::nn::Module {
    torch::nn::Sequential conv_layers1{ nullptr }, conv_layers2{ nullptr }, fc_layers{ nullptr };

    Net() {
        conv_layers1 = torch::nn::Sequential(
            // 第一层卷积：将 1 个输入通道映射为 32 个特征图，
            // kernel_size=3, stride=1, padding=1 保持输入尺寸不变（28×28），
            // 后接 BatchNorm 和 ReLU，再经过 2×2 的MaxPool降采样为14×14
            torch::nn::Conv2d(torch::nn::Conv2dOptions(1, 32, 3).padding(1)),
            torch::nn::BatchNorm2d(32),
            torch::nn::ReLU(),
            torch::nn::MaxPool2d(2)
        );
        conv_layers2 = torch::nn::Sequential(

            // 第二层卷积：将 32 个特征图映射到 64 个特征图，
            // 同样采用3×3卷积，padding=1(周围添加格子数量)，保证尺寸保持为14×14，然后经过2×2 MaxPool降采样为7×7
            torch::nn::Conv2d(torch::nn::Conv2dOptions(32, 64, 3).padding(1)),
            torch::nn::BatchNorm2d(64),
            torch::nn::ReLU(),
            torch::nn::MaxPool2d(2)
        );

        fc_layers = torch::nn::Sequential(
            // 将卷积层输出的 64 个 7×7 特征图展平，总特征数 64*7*7，
            // 经过全连接层映射到128个神经元，ReLU激活后再Dropout防止过拟合，
            // 最后映射到10个类别的输出
            torch::nn::Linear(64 * 7 * 7, 128),
            torch::nn::ReLU(),
            torch::nn::Dropout(0.5),
            torch::nn::Linear(128, 10)
        );

        // 注册模块
        register_module("conv_layers1", conv_layers1);
        register_module("conv_layers2", conv_layers2);
        register_module("fc_layers", fc_layers);
    }

    //c1 x in sizes:  [1, 1, 28, 28]
    //c1 x out sizes: [1, 32, 14, 14]
    //c2 x out sizes: [1, 64, 7, 7]
    //x view sizes:   [1, 3136]
    //x fc sizes:     [1, 10]
    
    torch::Tensor forward(torch::Tensor x) {
        x = conv_layers1->forward(x);

        //torch::nn::Sequential tt= conv_layers1.clone();
        //auto conv = dynamic_cast<torch::nn::Conv2dImpl*>(conv_layers1.get());
        //conv.to(torch::kCPU);
       // std::cout << "Weights of Conv2d Layer:\n" << conv->weight << std::endl;
        x = conv_layers2->forward(x);
        // 打印卷积层2的权重

        x = x.view({ x.size(0), -1 });
        x = fc_layers->forward(x);
        return x;
    }

};
*/
#endif  // MYNET_H_
